import fs from 'fs';
import path from 'path';
import { AICore } from "../AICore";
import { AIPanel } from "../AIPanel";
import { AIUtils } from "../AIUtils";
import { IDEAgent } from "../IDEAgent";
import { getImages, img2img } from "../Imagen";
import { AITask } from "./AITask";

class AITaskGenImageData{
    desc='一片叶子';
    negPrompt='';
    width=512;
    height=512;
    assetName='aigen.png'
    errorMsg:string;
    number=1;
    seed=0;
    tilex=false;
    tiley=false;
    generating=false;
    resultImgs:string[]=[];
    resultW:number;
    resultH:number;
    selected=0; //选中了哪张
    panorama=false;
    transparent=false;
    i2i=false;
    cfgScale=7;
    denoisingStrength=0.75;
}

class RangeItem{
    onChange:(value:number)=>void;
    title:gui.TextField;
    numinput:IEditor.NumericInput;
    slider:gui.Slider;
    constructor(item:gui.Widget, name:string|null, min:number, max:number, value:number, onChange:(value:number)=>void,callobj:any){
        this.title = item.getChild('title')
        this.numinput = item.getChild("numInput") as IEditor.NumericInput;
        this.slider = item.getChild("value") as gui.Slider;
        
        this.slider.value=value;
        this.numinput.text=value+'';
        name && (this.title.text=name);
        this.slider.min=min;
        this.slider.max=max;
        this.numinput.step=0.01;
        this.numinput.on('submit', () => {
            this.slider.value = this.numinput.value;
        });
        this.slider.on('changed',(e:gui.Event)=>{
            let v = (e.target as gui.Slider).value;
            this.numinput.text=v+'';
            if(onChange){
                onChange.call(callobj,v)
            }
        });

    }
}

export class AITaskGenImageNode extends AITask {
    data = new AITaskGenImageData()
    // 最终生成的文件
    protected _targetFile:string;
    resultAssetName:string=null;

    private _refImg:string=null;    //base64表示的
    private _maskImg:string=null;
    static EditObjName='图片生成';
    private onGenEnd:()=>void=()=>{};

    private onImgDragOver(e:Event){
        // 阻止默认行为，使得ondrop可以触发
        e.preventDefault();
        e.stopPropagation();
    }

    private async getItemString(item: DataTransferItem): Promise<string> {
        return new Promise<string>((res) => {
            item.getAsString(s => {
                res(s);
            });
        })
    }    
    private async getDropImageFile(event:any){
        let item = event.items[0];
        if (item.kind == 'string') {
            let s = await this.getItemString(item);
            if (s.startsWith(':')) {
                //这个是目录？
            }
            else {
                let fp: string;
                //文件
                if (s.startsWith('file:///')) {
                    fp = s.substring('file:///'.length);
                } else {
                    //具体的资源uuid
                    //await AICore.inst.memory.addFile(s, 'user');
                    //展示出来
                    let asset = await Editor.assetDb.getAsset(s);
                    fp = Editor.assetDb.getFullPath(asset);
                }
                //temp
                let ext = path.extname(fp).toLowerCase();
                if (ext == '.jpg' || ext == '.png') {
                    return fp;
                } 
            }
        } else if (item.kind == 'file') {
            let f = item.getAsFile();
            if(!f) return null;
            //await AICore.inst.memory.addFile(f, 'user');
            let type = f.type;
            if (type.startsWith('image')) {
                return (f as any).path;
            }
        }
        return null;
    }    

    private async _selectImage(index:number){
        let path = await IDEAgent.inst.download_img_from_url(this.data.resultImgs[index])
        let asset = await IDEAgent.inst.crateAIAssetFromFile(path, "image", this.data.assetName)
        Editor.scene?.setSelection([asset], "asset");
        this.data.selected = index;
        //保存ai信息
        IDEAgent.inst.setAIInfo(asset.id,this.data,true);
        this.genAsset = asset;
        this.resultAssetName = asset.fileName
        //this.renderUI();
    }

    async showUI(){
        let me = this;
        let data = this.data;
        let cont = await gui.UIPackage.createWidget("editorResources/AIUI/AIGenImage.widget");
        (cont as any).minWidth=520;
        cont.name='widget'; //aipanel根据这个来查找，并设置宽度
        let msgid = await AIPanel.ins.addMsg(cont,-1,{data:data});
        let gen_ = cont.getChild('generate') as gui.Box;
        let result_ = cont.getChild('result') as gui.List;
        let prompt_ = cont.getChild('prompt') as gui.TextField;
        let negprompt_ = cont.getChild('negprompt') as gui.TextField;
        let width_ = cont.getChild('width') as gui.TextInput;
        let height_ = cont.getChild('height') as gui.TextInput;
        let cont_ = cont.getChild('count') as gui.TextInput;
        let istransparent_ = cont.getChild('istransparent') as gui.Button;
        let istilex_ = cont.getChild('istilex') as gui.Button;
        let istiley_ = cont.getChild('istiley') as gui.Button;
        let i2i_ = cont.getChild('isI2I') as gui.Button;
        let i2i_cont_ = cont.getChild('I2ICont') as gui.Box;
        let cfgScale_ = cont.getChild('cfgScale') as gui.Widget;
        let denoisingStrength_ = cont.getChild('denoisingStrength') as gui.Widget;
        let refimg_ = i2i_cont_.getChild('refimg');
        let reftip_ = refimg_.getChild('label');

        let refimgshow = i2i_cont_.getChild('refimgshow') as gui.Loader;
        let close = cont.getChild('WindowCloseButton') as gui.Button;
        let select = cont.getChild('select') as gui.Shape;
        select.visible=false;
        //let shapedown_ = cont.getChild('Shapedown') as gui.Shape;

        prompt_.text = data.desc;
        prompt_.on('submit',(e:gui.Event)=>{ data.desc = (e.target as gui.TextField).text;})

        negprompt_.text = data.negPrompt;
        negprompt_.on('submit',(e:gui.Event)=>{data.negPrompt=(e.target as gui.TextField).text;});

        width_.text = data.width+'';
        width_.on('submit',(e:gui.Event)=>{
            data.width = parseInt((e.target as gui.TextField).text); 
            if(data.width>1024){
                data.width=1024;
                width_.text='1024';
                AICore.inst.tips('宽度最大为1024',1000);
            }
        })
        height_.text = data.height+'';
        height_.on('submit',(e:gui.Event)=>{
            data.height = parseInt((e.target as gui.TextField).text); 
            if(data.height>1024){
                data.height=1024;
                height_.text='1024';
                AICore.inst.tips('高度最大为1024',1000);
            }
        })

        cont_.text = data.number+'';
        cont_.on('submit',(e:gui.Event)=>{
            data.number = parseInt((e.target as gui.TextField).text); 
            if(data.number>4){
                data.number=4;
                cont_.text='4';
                AICore.inst.tips('个数最大为4',1000);
            }

        })
        
        istransparent_.selected = data.transparent;
        istransparent_.on('changed',(e:gui.Event)=>{ data.transparent = (e.target as gui.Button).selected;});

        istilex_.selected = data.tilex;
        istilex_.on('changed',(e:gui.Event)=>{ data.tilex = (e.target as gui.Button).selected;});

        istiley_.selected = data.tiley;
        istiley_.on('changed',(e:gui.Event)=>{ data.tiley = (e.target as gui.Button).selected;});

        i2i_.selected = data.i2i;
        function oni2i(b:boolean){
            i2i_cont_.visible=b;
            if(b){
                i2i_cont_.height=168;
            }else{
                i2i_cont_.height=0;
            }
            cont.height = result_.y+result_.height+10;
        }
        oni2i(data.i2i);

        i2i_.on('changed',()=>{
            data.i2i = i2i_.selected;
            oni2i(data.i2i);
            //i2i_cont_.removeSelf();
            //cont.height = gen_.y+gen_.height+50;
        })

        new RangeItem(cfgScale_,null,0,10, data.cfgScale, (v:number)=>{
            data.cfgScale = v;
        },this);

        new RangeItem(denoisingStrength_, null, 0,1, data.denoisingStrength, (v:number)=>{
            data.denoisingStrength=v;
        },this);

        function dragover(evt: gui.Event){
            var dt: DataTransfer = evt.data;
            dt.dropEffect = "copy";
            evt.preventDefault();
            //dragtip.visible=true;
        }
        function dragleave(evt:gui.Event){
            //dragtip.visible=false;
            evt.preventDefault();
        }

        function setRefImage(src:string){
            //console.log('文件名: ' + src);
            reftip_.visible=false;
            let filename = src.toLowerCase();
            let uriScheme = '';
            if (filename.endsWith('.png')) {
                uriScheme = 'data:image/png;base64,';
            } else if (filename.endsWith('.jpg') || filename.endsWith('.jpeg')) {
                uriScheme = 'data:image/jpg;base64,';
            } else if (filename.endsWith('.webp')) {
                uriScheme = 'data:image/webp;base64,';
            }

            let buff = fs.readFileSync(src);
            let base64 = buff.toString('base64')
            me._refImg = uriScheme+base64;

            refimgshow.src=new URL(src).toString();;
        }

        async function dropfunc(evt:gui.Event){
            var dt: DataTransfer = evt.data;
            dt.dropEffect = "copy";
            evt.preventDefault();
            //dragtip.visible=false;

            let r = await me.getDropImageFile(evt.data) as string;
            setRefImage(r);
        }

        //refimg_.on("native_dragover", dragover, null);
        cont.on("native_dragover", dragover, null);
        cont.on("native_drop", dropfunc, null);
        //refimgshow.on("native_drop", dropfunc, null);
        //refimg_.on('native_dragleave',dragleave,null);

        this.onGenEnd=()=>{
            //删除
            for(let i=result_.numChildren; i>data.resultImgs.length; i--){
                result_.removeChildAt(i-1);
            }
            //添加
            for(let i=result_.numChildren; i<data.resultImgs.length; i++){
                let r1 = new gui.Loader();
                result_.addChild(r1);
                r1.on('right_click', (e:gui.Event) => {
                    let menu = IEditor.Menu.create([ 
                        { 
                            label: "作为参考图" , 
                            click : async function() { 
                                i2i_.selected = true;
                                data.i2i = true;
                                oni2i(true);
                                let p = await IDEAgent.inst.download_img_from_url(r1.src,true)
                                setRefImage(p);
                            }
                        }
                    ]);
                    //当需要弹出时
                    menu.show();
                    e.preventDefault();
                    //阻止继续冒泡。否则会触发panel的右键导致不显示菜单
                    e.stopPropagation();
                });        
                r1.on('click',(e:gui.Event)=>{
                    me._selectImage(i);
                    select.visible=true;
                    select.removeSelf();
                    r1.addChild(select);
                    select.x = 0;
                    select.y = 0;
                    select.width=r1.width;
                    select.height=r1.height;
                });
            }

            for(let i=0; i<data.resultImgs.length; i++){
                let r = result_.getChildAt(i) as gui.Loader;
                r.width=200;
                r.height =200*(data.resultH/data.resultW); 
                r.src = data.resultImgs[i];
            }
        }

        result_.on('size_changed',()=>{
            cont.height = result_.y+result_.height+10;
        });

        gen_.on('click',()=>{
            data.seed = Math.floor(Math.random() * 1000000);
            this.genImage(false);
        })

        close.on('click', (e:gui.Event)=>{
            AIPanel.ins.removeMsg(msgid);
        }, this);

        return cont;
    }

    /// 执行任务 ///
    async genImage(reuse:boolean = true){
        //检查是否有同名文件
        this.data.assetName='img_'+Date.now()+'.png';
        this.data.generating=true;
        let wmsg = await AIPanel.ins.startWaitMsg('生成中');
        let data = this.data;
        data.resultW=data.width;
        data.resultH=data.height;
        let tilex = data.tilex;
        let tiley = data.tiley;
        let panorama = data.panorama;
        let prompt = `${data.desc}`;
        let negPrompt = data.negPrompt;
        let translate = false
        if(translate){
            prompt = await AIUtils.translation(prompt);
        }
        if(data.negPrompt){
            negPrompt = await AIUtils.translation(data.negPrompt);
        }
        if(data.i2i && this._refImg){
            prompt = await AIUtils.translation(prompt); //图生图的服务器没有翻译。所以这里先翻译一下
            let urls = await img2img(this._refImg,
                    this._maskImg,prompt, {
                        alpha:data.transparent,
                        batch_size: Math.min(data.number, 4),
                        seed:data.seed,
                        width:data.width,
                        height:data.height,
                        cfg_scale:data.cfgScale,
                        denoising_strength:data.denoisingStrength,
                    }, data.tilex, data.tiley);

            if(!urls){
                AIPanel.ins.removeMsg(wmsg);
                AIPanel.ins.addMsg('图片生成失败！',-2);
                this.data.generating=false;
                return null;
            }

            if(this._isCancel || !urls){
                AIPanel.ins.removeMsg(wmsg);
                this.data.generating=false;
                return null;
            }
            //得到结果,结果是一个图片资源，保存起来
            let imgfiles: string[] = urls;
            this.data.resultImgs = imgfiles;
            this.data.generating = false;
            //aiPanel.removeMsg(tipmsgid);
            AIPanel.ins.removeMsg(wmsg);
            await this._selectImage(0);
            AICore.inst.tips('图片已经生成');
            //if(!showUI){
            //    AIPanel.ins.addMsg( `<img src="${Editor.assetsPath+"/aigc/image/"+this.resultAssetName}" width=${200} height=${200} />`,-1);
            //}
            this.onGenEnd();
            return this.genAsset;
        }else{
            let urls = await getImages({
                    prompt,
                    width: data.width,
                    height: data.height,
                    batch_size: Math.min(data.number, 4),
                    seed:data.seed,
                    negative_prompt:'nsfw '+negPrompt,
                    reuse: reuse,
                    trans: true
                     }, data.transparent, tilex, tiley, panorama)
            if(!urls){
                AIPanel.ins.removeMsg(wmsg);
                AIPanel.ins.addMsg('图片生成失败！',-2);
                this.data.generating=false;
                return null;        
            }
            if(this._isCancel || !urls){
                AIPanel.ins.removeMsg(wmsg);
                this.data.generating=false;
                return null;
            }
            let imgfiles: string[] = urls;//[];
            this.data.resultImgs = imgfiles;
            this.data.generating = false;
            AIPanel.ins.removeMsg(wmsg);
            await this._selectImage(0);
            AICore.inst.tips('图片已经生成');
            //if(!showUI){
            //    AIPanel.ins.addMsg( `<img src="${Editor.assetsPath+"/aigc/image/"+this.resultAssetName}" width=${200} height=${200} />`,-1);
            //}
            this.onGenEnd();
            return this.genAsset;
        }
    }

    async execute(){
        this.data.seed = Math.floor(Math.random() * 1000000);
        this.showUI();
        // 等待界面交互，直到用户点击ok或者说好了
        await this.waitIt();
        this.resultAssetName = this.data.assetName;
        let asset = await Editor.assetDb.getAsset(this.resultAssetName);
        this.genAsset = asset;
    }
}

